#!/usr/bin/env python3
"""
Harvest sequence and element definitions directly from BUFR messages using pybufrkit.

This is useful for NESDIS radiance files that embed local Table-D sequences (e.g.,
3-61-234, 3-10-198).  The script walks the message templates, extracts the target
sequences (and any nested dependencies), and emits ecCodes-compatible `sequence.def`
and `element.table` files.
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Dict, Iterable, List, Set, Tuple

from pybufrkit.commands import generate_bufr_message
from pybufrkit.decoder import Decoder
from pybufrkit.descriptors import FixedReplicationDescriptor, SequenceDescriptor


def descriptor_id(desc) -> int:
    return desc.id


def flatten_sequence(
    node,
    target_sequences: Set[int],
    sequence_defs: Dict[int, List[int]],
    element_meta: Dict[int, object],
    visited: Set[int],
) -> List[int]:
    seq_id = descriptor_id(node.descriptor)
    if seq_id in visited:
        return sequence_defs[seq_id]

    entries: List[int] = []
    for child in node.members:
        desc = child.descriptor
        fid = descriptor_id(desc)
        F = desc.F

        if F == 3:
            entries.append(fid)
            flatten_sequence(child, target_sequences, sequence_defs, element_meta, visited)
        elif F == 1 and isinstance(desc, FixedReplicationDescriptor):
            entries.append(fid)
            if child.members:
                first = child.members[0]
                flatten_sequence(first, target_sequences, sequence_defs, element_meta, visited)
                entries.append(descriptor_id(first.descriptor))
        else:
            entries.append(fid)
            if F == 0:
                element_meta.setdefault(fid, desc)

    sequence_defs[seq_id] = entries
    visited.add(seq_id)
    return entries


def harvest_sequences(
    bufr_file: Path,
    sequence_ids: Set[int],
) -> Tuple[Dict[int, List[int]], Dict[int, object]]:
    decoder = Decoder(fallback_or_ignore_missing_tables=True)
    with bufr_file.open("rb") as fh:
        data = fh.read()

    sequence_defs: Dict[int, List[int]] = {}
    element_meta: Dict[int, object] = {}
    visited: Set[int] = set()
    seen: Set[int] = set()

    for message in generate_bufr_message(
        decoder, data, file_path=str(bufr_file), wire_template_data=True
    ):
        template = message.sections[3].template_data.value
        template.wire()
        subsets = template.decoded_nodes_all_subsets
        if not subsets:
            continue
        for node in subsets[0]:
            desc = node.descriptor
            if isinstance(desc, SequenceDescriptor):
                sid = descriptor_id(desc)
                if sid in sequence_ids or sid not in visited:
                    flatten_sequence(node, sequence_ids, sequence_defs, element_meta, visited)
                    seen.add(sid)
        if sequence_ids.issubset(seen):
            break

    missing = sequence_ids.difference(seen)
    if missing:
        raise SystemExit(f"Sequences {sorted(missing)} not found in {bufr_file}")

    return sequence_defs, element_meta


def write_sequence_def(sequence_map: Dict[int, List[int]], output_path: Path) -> None:
    lines = [
        "# Auto-generated from BUFR file via harvest_bufr_sequences.py",
    ]
    for code in sorted(sequence_map):
        values = [f"{val:06d}" for val in sequence_map[code]]
        chunks: List[str] = []
        line: List[str] = []
        for idx, val in enumerate(values, start=1):
            line.append(val)
            if idx % 6 == 0 and idx != len(values):
                chunks.append(", ".join(line))
                line = []
        if line:
            chunks.append(", ".join(line))
        if len(chunks) == 1:
            body = chunks[0]
        else:
            body = ",\n              ".join(chunks)
        lines.append(f'"{code:06d}" = [  {body} ]\n')

    output_path.write_text("\n".join(lines), encoding="ascii")


def write_element_table(element_map: Dict[int, object], extra_elements: Iterable[int], output_path: Path) -> None:
    lines = [
        "#code|abbreviation|type|name|unit|scale|reference|width|crex_unit|crex_scale|crex_width\n"
    ]
    emitted: Set[int] = set()

    for desc_id in sorted(element_map):
        desc = element_map[desc_id]
        abbreviation = desc.name.strip().lower().replace(" ", "_")
        unit = desc.unit.strip() if isinstance(desc.unit, str) else ""
        lines.append(
            f"{desc_id:06d}|{abbreviation}|double|{desc.name.strip()}|{unit}|{desc.scale}|{desc.refval}|"
            f"{desc.nbits}|{unit or 'NA'}|0|0\n"
        )
        emitted.add(desc_id)

    for extra in extra_elements:
        if extra in emitted:
            continue
        lines.append(f"{extra:06d}|descriptor_{extra}|double|LOCAL DESCRIPTOR {extra}|NA|0|0|16|NA|0|0\n")

    output_path.write_text("".join(lines), encoding="ascii")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("bufr_file", type=Path, help="Input BUFR file")
    parser.add_argument("output_dir", type=Path, help="Directory to write sequence.def/element.table")
    parser.add_argument(
        "--sequence",
        type=int,
        action="append",
        required=True,
        help="Sequence descriptor (e.g., 361234). Can be provided multiple times.",
    )
    parser.add_argument(
        "--extra-element",
        type=int,
        action="append",
        default=[],
        help="Optional Table-B descriptor to append to element.table.",
    )
    args = parser.parse_args()

    target_sequences = set(args.sequence)
    output_dir = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)

    sequence_defs, element_meta = harvest_sequences(args.bufr_file, target_sequences)

    seq_path = output_dir / "sequence.def"
    write_sequence_def(sequence_defs, seq_path)

    elem_path = output_dir / "element.table"
    write_element_table(element_meta, args.extra_element, elem_path)

    print(f"Wrote {seq_path}")
    print(f"Wrote {elem_path}")


if __name__ == "__main__":
    main()
